package com.maweirdos.ai.core.adviser;

import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.ai.chat.prompt.Prompt;
import reactor.core.publisher.Flux;

/**
 * 自定义重读顾问器
 */
public class ReReadingadvisor implements CallAdvisor, StreamAdvisor {

    public static final String RE_READ_TEXT = "Read the question again: ";

    private final int order;

    public ReReadingadvisor(int order) {
        this.order = order;
    }

    /**
     * 执行请求前，改写 Prompt
     * @param request
     * @return
     */
    private ChatClientRequest before(ChatClientRequest request) {

        // 使用augmentUserMessage（增强用户消息）方法，在用户消息后面添加重读提示
        Prompt rePrompt = request.prompt().augmentUserMessage(userMessage -> userMessage.mutate()
                .text(userMessage.getText() + System.lineSeparator() + RE_READ_TEXT + userMessage.getText())
                .build()
        );

        // 创建新的请求
        return ChatClientRequest.builder()
                .prompt(rePrompt)
                .context(request.context())
                .build();
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
        return callAdvisorChain.nextCall(this.before(chatClientRequest));
    }

    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain streamAdvisorChain) {
        return streamAdvisorChain.nextStream(this.before(chatClientRequest));
    }

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        return this.order;
    }

    public static Builder builder() {
        return new Builder();
    }

    public static final class Builder {

        private int order = 0;

        private Builder() {

        }

        public Builder order(int order) {
            this.order = order;
            return this;
        }

        public ReReadingadvisor build() {
            return new ReReadingadvisor(this.order);
        }
    }
}
